package com.jboltai.event.ai;

import cn.hutool.core.util.StrUtil;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.config.JBoltAIConfig;
import com.jboltai.event.Event;
import com.jboltai.event.EventType;
import com.jboltai.event.handler.EventHandler;
import com.jboltai.resource.ai.AIModel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public abstract class AIEvent<T extends Event, ResultType,HandlerType extends EventHandler> extends Event<T, ResultType,HandlerType> implements IAIEvent<T, ResultType,HandlerType> {

    private static final Logger LOGGER = LoggerFactory.getLogger(JBoltAIConfig.LOGGER_NAME);

    private static final String TOP_K_KEY = "top_k";

    /**
     * 创建一个ai事件
     */
    protected AIEvent(EventType eventType, AIModel model, String modelName) {
        super(eventType);
        this.modelName = modelName;
        this.model = model;

        if (JBoltAIConfig.getMaxAIEventWaitTimeout() > 0) {
            this.waitTimestamp = System.currentTimeMillis() + JBoltAIConfig.getMaxAIEventWaitTimeout();
        } else {
            this.waitTimestamp = 0;
        }

    }


    /**
     * 事件执行时使用的大模型类型
     */
    protected AIModel model;

    /**
     * 具体使用的大模型
     */
    protected String modelName;

    /**
     * 对大模型的system角色设置
     */
    protected String systemMessage;

    /**
     * 上下文消息
     */
    protected List<AIMessage> contextMessages;

    /**
     * 最后一条用户消息指令，保证插到消息列表的最后
     */
    protected String prompt;

    /**
     * 最后一条用户消息指令要带的图片
     */
    protected String[] promptImgs;

    /**
     * 输入内容是否包含文字
     */
    protected boolean textInput = true;

    /**
     * 输入内容是否包含图片
     */
    protected boolean imageInput;

    /**
     * 输入内容是否包含音频
     */
    protected boolean audioInput;

    /**
     * 输入内容是否包含视频
     */
    protected boolean videoInput;

    /**
     * 开启思考模式
     */
    protected Boolean thinkingEnable;

    /**
     * 指定用哪个resourceId执行任务，null为自动分配
     */
    protected String aiResourceId;


    /**
     * 温度值，越低回答越精确，越高回答越随机。
     */
    protected Double temperature;
    /**
     * 多样性 top_p
     */
    protected Double topP;

    /**
     * 大模型请求需要的其他参数
     */
    protected Map<String,Object> params;

    @Override
    public Map<String,Object> getParams() {
        return params;
    }

    @Override
    public T setParams(Map<String, Object> params) {
        this.params = params;
        return (T) this;
    }

    @Override
    public T addParam(String key, Object value){
        if(this.params==null){
            this.params = new HashMap<>();
        }
        this.params.put(key,value);
        return (T) this;
    }

    @Override
    public <V> V getParam(String key) {
        if (this.params == null) {
            return null;
        }
        return (V) this.params.get(key);
    }

    /**
     * 将该事件分配到指定的资源上
     * @param resourceId
     * @return
     */
    @Override
    public T bindAIResource(String resourceId) {
        this.aiResourceId = resourceId;
        return (T) this;
    }


    @Override
    public T setTemperature(Double temperature) {
        this.temperature = temperature;
        return (T) this;
    }

    @Override
    public T setTopP(Double topP) {
        if (topP != null && (topP > 1 || topP < 0)) {
            throw new IllegalArgumentException("top_p的参数范围是0-1");
        }
        this.topP = topP;
        return (T) this;
    }
    @Override
    public T setTopK(Integer topK) {
        this.addParam(TOP_K_KEY, topK);
        return (T) this;
    }

    @Override
    public Integer getTopK() {
        return this.getParam(TOP_K_KEY);
    }



    @Override
    public String getAIResource() {
        return aiResourceId;
    }





    @Override
    public String getModelName() {
        return modelName;
    }

    @Override
    public AIModel getModel() {
        return model;
    }


    @Override
    public Double getTemperature() {
        return temperature;
    }

    @Override
    public Double getTopP() {
        return topP;
    }


    /**
     * 获取所有的消息列表，如果设置了prompt，会将该消息插到消息列表的最后返回
     * @return
     */
    @Override
    public List<AIMessage> getMessages() {
        List<AIMessage> messages = new ArrayList<>();
        if(StrUtil.isNotBlank(systemMessage)){
            messages.add(0,AIMessage.ofSystem(systemMessage));
        }
        if (this.contextMessages != null && !this.contextMessages.isEmpty()) {
            messages.addAll(contextMessages);
        }
        if(StrUtil.isNotBlank(prompt)){
            AIMessage msg = AIMessage.of(prompt);
            if (promptImgs != null && promptImgs.length != 0) {
                msg.setImgUrls(promptImgs);
            }
            messages.add(msg);
        }
        return messages;
    }

    @Override
    public T contextMessages(List<AIMessage> contextMessages){
        this.contextMessages = contextMessages;
        return (T) this;
    }


    /**
     * 添加system角色消息
     * @param content
     * @return
     */
    @Override
    public T systemMessage(String content) {
        this.systemMessage = content;
        return (T) this;
    }
    /**
     * 添加user角色消息
     * @param content
     * @return
     */
    @Override
    public T userMessage(String content) {
        if (this.contextMessages == null) {
            this.contextMessages = new ArrayList<>();
        }
        this.contextMessages.add(AIMessage.of(content));
        return (T) this;
    }
    /**
     * 添加user角色消息
     * @param content
     * @param imgUrls
     * @return
     */
    @Override
    public T userMessage(String content, String... imgUrls) {
        if (this.contextMessages == null) {
            this.contextMessages = new ArrayList<>();
        }
        this.contextMessages.add(AIMessage.of(content, imgUrls));
        if (imgUrls != null && imgUrls.length != 0) {
            this.imageInput = true;
        }
        return (T) this;
    }
    /**
     * 添加Assistant角色消息
     * @param content
     * @return
     */
    @Override
    public T assistantMessage(String content) {
        if (this.contextMessages == null) {
            this.contextMessages = new ArrayList<>();
        }
        this.contextMessages.add(AIMessage.ofAssistant(content));
        return (T) this;
    }


    @Override
    public T setModel(AIModel model) {
        this.model = model;
        return (T) this;
    }

    /**
     * 指定使用的大模型，请注意，除非您确认该模型一定可以成功匹配注册的ai资源，否则不要自行调用该方法
     * @param modelName
     * @return
     */
    @Override
    public T setModelName(String modelName) {
        this.modelName = modelName;
        return (T) this;
    }


    @Override
    public String getPrompt() {
        return prompt;
    }

    @Override
    public T prompt(String prompt) {
        this.prompt = prompt;
        return (T) this;
    }

    /**
     * 最后一条用户消息指令要带的图片。
     * 不同的大模型，对于图片支持的格式不同。比如chatgpt4下只支持传入图片url，ollama下支持url和base64两种格式。<br/>
     * 请根据不同大模型进行传参。
     * @param imgs
     * @return
     */
    @Override
    public T imgs(String... imgs) {
        if (imgs != null && imgs.length != 0) {
            this.imageInput = true;
        }
        this.promptImgs = imgs;
        return (T) this;
    }


    /**
     * 通过模版构造prompt，占位符使用{}表示
     * @param template
     * @param values
     * @return
     */
    @Override
    public T prompt(String template, Object... values) {
        this.prompt = StrUtil.format(template, values);
        return (T) this;
    }

    @Override
    public boolean isTextInput() {
        return textInput;
    }

    @Override
    public void setTextInput(boolean textInput) {
        this.textInput = textInput;
    }

    @Override
    public boolean isImageInput() {
        return imageInput;
    }

    @Override
    public void setImageInput(boolean imageInput) {
        this.imageInput = imageInput;
    }

    @Override
    public boolean isAudioInput() {
        return audioInput;
    }

    @Override
    public void setAudioInput(boolean audioInput) {
        this.audioInput = audioInput;
    }

    @Override
    public boolean isVideoInput() {
        return videoInput;
    }

    @Override
    public void setVideoInput(boolean videoInput) {
        this.videoInput = videoInput;
    }

    public List<AIMessage> getContextMessages() {
        return contextMessages;
    }

    public Boolean isThinkingEnable() {
        return thinkingEnable;
    }

    public T  setThinkingEnable(Boolean thinkingEnable) {
        this.thinkingEnable = thinkingEnable;
        return (T) this;
    }
}
